#!/usr/bin/env python3
"""
4波段配置集成测试脚本
测试从数据加载到模型预测的完整流程
"""

import torch
import numpy as np
import os
import sys
from datetime import datetime, timedelta
import pandas as pd

# 添加项目路径
sys.path.append(os.path.dirname(__file__))

from data_utils import WeatherDataset, load_config
from models import CombinedModel

def test_data_loading():
    """测试数据加载功能"""
    print("=" * 50)
    print("测试数据加载功能")
    print("=" * 50)
    
    try:
        # 加载配置
        config = load_config()
        print(f"✅ 配置加载成功")
        print(f"   波段配置: {config['data']['himawari_bands']}")
        print(f"   输入通道数: {config['model']['num_channels']}")
        
        # 创建数据集
        dataset = WeatherDataset(
            station_csv=config['data']['station_csv'],
            himawari_dir=config['data']['himawari_dir'],
            historical_days=config['data']['historical_days'],
            future_hours=config['data']['future_hours']
        )
        
        print(f"✅ 数据集创建成功")
        print(f"   样本数量: {len(dataset)}")
        
        # 测试一个样本
        if len(dataset) > 0:
            sample = dataset[0]
            print(f"✅ 样本加载成功")
            print(f"   station_data shape: {sample['station_data'].shape}")
            print(f"   himawari_data shape: {sample['himawari_data'].shape}")
            print(f"   future_data shape: {sample['future_data'].shape}")
            
            # 检查波段数量
            expected_channels = len(config['data']['himawari_bands'])
            actual_channels = sample['himawari_data'].shape[0]
            if actual_channels == expected_channels:
                print(f"✅ 波段数量正确: {actual_channels}")
            else:
                print(f"❌ 波段数量错误: 期望{expected_channels}, 实际{actual_channels}")
                return False
                
        return True
        
    except Exception as e:
        print(f"❌ 数据加载测试失败: {e}")
        return False

def test_model_creation():
    """测试模型创建功能"""
    print("\n" + "=" * 50)
    print("测试模型创建功能")
    print("=" * 50)
    
    try:
        # 加载配置
        config = load_config()
        
        # 创建模型
        model = CombinedModel(config)
        print(f"✅ 模型创建成功")
        
        # 检查参数数量
        total_params = sum(p.numel() for p in model.parameters())
        trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
        print(f"   总参数: {total_params:,}")
        print(f"   可训练参数: {trainable_params:,}")
        
        # 测试前向传播
        batch_size = 2
        historical_points = config['data']['historical_days'] * 24 * 4
        station_data = torch.randn(batch_size, historical_points, 3)
        himawari_data = torch.randn(batch_size, config['model']['num_channels'], 224, 224)
        
        with torch.no_grad():
            output = model(station_data, himawari_data)
        
        print(f"✅ 前向传播成功")
        print(f"   输入 station_data: {station_data.shape}")
        print(f"   输入 himawari_data: {himawari_data.shape}")
        print(f"   输出: {output.shape}")
        
        # 检查输出维度
        expected_future_points = config['data']['future_hours'] * 4
        expected_output_vars = config['model']['num_output_vars']
        if output.shape == (batch_size, expected_future_points, expected_output_vars):
            print(f"✅ 输出维度正确")
        else:
            print(f"❌ 输出维度错误")
            return False
            
        return True
        
    except Exception as e:
        print(f"❌ 模型创建测试失败: {e}")
        return False

def test_training_compatibility():
    """测试训练兼容性"""
    print("\n" + "=" * 50)
    print("测试训练兼容性")
    print("=" * 50)
    
    try:
        # 加载配置
        config = load_config()
        
        # 创建数据加载器
        from torch.utils.data import DataLoader
        dataset = WeatherDataset(
            station_csv=config['data']['station_csv'],
            himawari_dir=config['data']['himawari_dir'],
            historical_days=config['data']['historical_days'],
            future_hours=config['data']['future_hours']
        )
        
        if len(dataset) > 0:
            # 创建数据加载器
            dataloader = DataLoader(dataset, batch_size=2, shuffle=False)
            
            # 获取一个批次
            batch = next(iter(dataloader))
            
            print(f"✅ 数据加载器创建成功")
            print(f"   批次 station_data: {batch['station_data'].shape}")
            print(f"   批次 himawari_data: {batch['himawari_data'].shape}")
            print(f"   批次 future_data: {batch['future_data'].shape}")
            
            # 检查批次维度
            if batch['himawari_data'].shape[1] == config['model']['num_channels']:
                print(f"✅ 批次波段数量正确")
            else:
                print(f"❌ 批次波段数量错误")
                return False
                
        return True
        
    except Exception as e:
        print(f"❌ 训练兼容性测试失败: {e}")
        return False

def test_prediction_compatibility():
    """测试预测兼容性"""
    print("\n" + "=" * 50)
    print("测试预测兼容性")
    print("=" * 50)
    
    try:
        # 加载配置
        config = load_config()
        
        # 创建模型
        model = CombinedModel(config)
        
        # 模拟预测输入
        batch_size = 1
        historical_points = config['data']['historical_days'] * 24 * 4
        station_data = torch.randn(batch_size, historical_points, 3)
        himawari_data = torch.randn(batch_size, config['model']['num_channels'], 224, 224)
        
        # 预测模式
        model.eval()
        with torch.no_grad():
            predictions = model(station_data, himawari_data)
        
        print(f"✅ 预测模式测试成功")
        print(f"   预测输出: {predictions.shape}")
        
        # 检查预测输出
        expected_future_points = config['data']['future_hours'] * 4
        expected_output_vars = config['model']['num_output_vars']
        if predictions.shape == (batch_size, expected_future_points, expected_output_vars):
            print(f"✅ 预测输出维度正确")
        else:
            print(f"❌ 预测输出维度错误")
            return False
            
        return True
        
    except Exception as e:
        print(f"❌ 预测兼容性测试失败: {e}")
        return False

def main():
    """主测试函数"""
    print("开始4波段配置集成测试")
    print("=" * 60)
    
    tests = [
        ("数据加载", test_data_loading),
        ("模型创建", test_model_creation),
        ("训练兼容性", test_training_compatibility),
        ("预测兼容性", test_prediction_compatibility)
    ]
    
    passed_tests = 0
    total_tests = len(tests)
    
    for test_name, test_func in tests:
        try:
            if test_func():
                print(f"✅ {test_name}测试通过\n")
                passed_tests += 1
            else:
                print(f"❌ {test_name}测试失败\n")
        except Exception as e:
            print(f"❌ {test_name}测试异常: {e}\n")
    
    print("=" * 60)
    print(f"测试结果: {passed_tests}/{total_tests} 通过")
    
    if passed_tests == total_tests:
        print("🎉 所有测试通过！4波段配置完整可用")
        return True
    else:
        print("⚠️  部分测试失败，需要进一步检查")
        return False

if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)
